'''
    module_of_msg.py
    author: yizimi-yuanxin
    update: 23.3.15
'''

from data_type_extensible import *


def get_string_mode():
    '''
    创建string类型的模式，仅适用于data
    '''
    string_mode = msg_mode(0x00000000, "string")
    string_mode.add_fixed_mode_seg(0, "data")
    return string_mode

def get_pcapheader_mode():
    '''
    创建pcap header的格式，用于读取pcap文件
    '''
    Pcapheader_mode = msg_mode(0xFFFFFFFF, "Pcapheader")
    Pcapheader_mode.add_fixed_mode_seg(4 * 8, 'Magic')
    Pcapheader_mode.add_fixed_mode_seg(2 * 8, 'Major')
    Pcapheader_mode.add_fixed_mode_seg(2 * 8, 'Minor')
    Pcapheader_mode.add_fixed_mode_seg(4 * 8, 'ThisZone')
    Pcapheader_mode.add_fixed_mode_seg(4 * 8, 'SigFigs')
    Pcapheader_mode.add_fixed_mode_seg(4 * 8, 'SnapLen')
    Pcapheader_mode.add_fixed_mode_seg(4 * 8, 'LinkType')
    Pcapheader_mode.set_data_mode(-1)
    return Pcapheader_mode

def get_packetheader_mode():
    '''
    创建packet header的格式，用于pcap文件分隔报文和确定报文长度
    '''
    Packetheader_mode = msg_mode(0xFFFFFFFE, "Packetheader")
    Packetheader_mode.add_fixed_mode_seg(4 * 8, "ts_sec")
    Packetheader_mode.add_fixed_mode_seg(4 * 8, "ts_usec")
    Packetheader_mode.add_fixed_mode_seg(4 * 8, "incl_len")
    Packetheader_mode.add_fixed_mode_seg(4 * 8, "orig_len")
    Packetheader_mode.add_fixed_mode_seg(-1, "data")
    Packetheader_mode.set_datalen_name("incl_len")
    Packetheader_mode.set_data_mode(-1)
    next_list = {
        0: "BSD loopback", 			# undefined
        1: "Ethernet_II",
        6: "802.5 Token Ring",		# undefined
        7: "ARCnet",				# undefined
        8: "SLIP", 					# undefined
        9: "PPP",					# undefined
        10: "FDDI", 				# undefined
        105: "802.11"				# undefined
    }
    Packetheader_mode.set_next_mode_code("pcap_header")
    Packetheader_mode.set_next_mode_list(next_list)
    # There is some questions.....
    # TODO: A part of Pcapheader need to use for Packetheader's analysis
    return Packetheader_mode

def get_Ethernet_II_mode():
    '''
    
    '''
    Ethernet_II_mode = msg_mode(0x00000001, "Ethernet_II")
    Ethernet_II_mode.add_fixed_mode_seg(6 * 8, "dMAC")
    Ethernet_II_mode.add_fixed_mode_seg(6 * 8, "sMAC")
    Ethernet_II_mode.add_fixed_mode_seg(2 * 8, "Type")
    Ethernet_II_mode.add_fixed_mode_seg(-1, "data")
    Ethernet_II_mode.set_datalen_name(None)
    Ethernet_II_mode.set_data_mode(0)
    next_list = {
        0x0800: "IPv4",
        0x0803: "ECMA Internet", 		# undefined
        0x0806: "ARP",					# undefined
        0x0808: "Frame Relay ARP", 		# undefined
        0x8035: "Reverse ARP",			# undefined
        0x8046: "AT&T",					# undefined
        0x8047: "AT&T",					# undefined
        0x876B: "TCP/IP Compression", 	# undefined
        0x86DD: "IPv6", 				# undefined
        0x880B: "PPP"					# undefined
    }
    Ethernet_II_mode.set_next_mode_code("Type")
    Ethernet_II_mode.set_next_mode_list(next_list)
    return Ethernet_II_mode

def get_IPv4_mode():
    '''
    
    '''
    IPv4_mode = msg_mode(0x00000002, "IPv4")
    IPv4_mode.add_fixed_mode_seg(4, "version")
    IPv4_mode.add_fixed_mode_seg(4, "headerLen")
    IPv4_mode.add_fixed_mode_seg(8, "tos")
    IPv4_mode.add_fixed_mode_seg(16, "packetLen")
    IPv4_mode.add_fixed_mode_seg(16, "protoId")
    IPv4_mode.add_fixed_mode_seg(3, "flags")
    IPv4_mode.add_fixed_mode_seg(13, "frag")
    IPv4_mode.add_fixed_mode_seg(8, "ttl")
    IPv4_mode.add_fixed_mode_seg(8, "proto")
    IPv4_mode.add_fixed_mode_seg(16, "chksum")
    IPv4_mode.add_fixed_mode_seg(32, "srcIp")
    IPv4_mode.add_fixed_mode_seg(32, "dstIp")
    IPv4_mode.set_data_mode(0)
    
    def optfunc(msg: msg_struct):
        return (msg.get_seg("headerLen").seg_data * 4 - 20) * 8
    IPv4_mode.add_variable_mode_seg(optfunc, "options")
    
    def datafunc(msg: msg_struct):
        return (msg.get_seg("packetLen").seg_data - msg.get_seg("headerLen").seg_data * 4) * 8
    IPv4_mode.add_variable_mode_seg(datafunc, "data")
    # IPv4_mode.add_fixed_mode_seg(-1, "data")
    # IPv4_mode.set_datalen_name("packetLen")
    next_list = {
        1: "ICMP", 						# undefined
        2: "IGMP",						# undefined
        4: "IP", 						# undefined
        6: "TCP", 						
        8: "EGP", 						# undefined
        17: "UDP",						# undefined
        41: "IPv6", 					# undefined
        89: "OSPF"						# undefined
    }
    IPv4_mode.set_next_mode_code("proto")
    IPv4_mode.set_next_mode_list(next_list)
    return IPv4_mode


def get_TCP_mode():
    '''
    
    '''
    TCP_mode = msg_mode(0x00000003, "TCP")
    TCP_mode.add_fixed_mode_seg(16, "sport")
    TCP_mode.add_fixed_mode_seg(16, "dport")
    TCP_mode.add_fixed_mode_seg(32, "seq")
    TCP_mode.add_fixed_mode_seg(32, "ack")
    TCP_mode.add_fixed_mode_seg(4, "dataofs")
    TCP_mode.add_fixed_mode_seg(4, "reserved")
    TCP_mode.add_fixed_mode_seg(8, "flags")
    TCP_mode.add_fixed_mode_seg(16, "window")
    TCP_mode.add_fixed_mode_seg(16, "chksum")
    TCP_mode.add_fixed_mode_seg(16, "urgptr")
    
    def optfunc(msg: msg_struct):
        return (msg.get_seg("dataofs").seg_data * 4 - 20) * 8
    TCP_mode.add_variable_mode_seg(optfunc, "options")
    
    TCP_mode.add_fixed_mode_seg(-1, "data")
    TCP_mode.set_data_mode(0)
    TCP_mode.set_datalen_name(None)
    
    # def datafunc(msg: msg_struct):
    #     return (msg.get_seg("packetLen").seg_data - msg.get_seg("headerLen").seg_data * 4) * 8
    # IPv4_mode.add_variable_mode_seg(datafunc, "data")
    
    TCP_mode.set_next_mode_code("-1")
    return TCP_mode
    
def get_UDP_mode():
    '''
    
    '''
    UDP_mode = msg_mode(0x00000004, "UDP")
    UDP_mode.add_fixed_mode_seg(16, "sport")
    UDP_mode.add_fixed_mode_seg(16, "dport")
    UDP_mode.add_fixed_mode_seg(16, "packetLen")
    UDP_mode.add_fixed_mode_seg(16, "chksum")
    
    def datafunc(msg: msg_struct):
        return (msg.get_seg("packetLen").seg_data - 8) * 8
    UDP_mode.add_variable_mode_seg(datafunc, "data")
    UDP_mode.set_next_mode_code("-1")
    return UDP_mode

str_mode = get_string_mode()
EthII_mode = get_Ethernet_II_mode()
Pcaph_mode = get_pcapheader_mode()
Pkt_mode = get_packetheader_mode()
IPv4_mode = get_IPv4_mode()
TCP_mode = get_TCP_mode()
UDP_mode = get_UDP_mode()

# rev_exc_dict = {
#     0x00000000: get_string_mode(),
#     0xFFFFFFFF: get_pcapheader_mode(),
#     0xFFFFFFFE: get_packetheader_mode(),
#     0x00000001: get_Ethernet_II_mode(),
#     0x00000002: get_IPv4_mode(),
#     0x00000003: get_TCP_mode()
# }

def rev_exc_dict(index: int):
    if index == 0x00000000: 
        return str_mode
    if index == 0xFFFFFFFF: 
        return Pcaph_mode
    if index == 0xFFFFFFFE: 
        return Pkt_mode
    if index == 0x00000001: 
        return EthII_mode
    if index == 0x00000002: 
        return IPv4_mode
    if index == 0x00000003: 
        return TCP_mode
    if index == 0x00000004:
        return UDP_mode
    
    return str_mode

# rev_exc_dict = {
#     0x00000000: str_mode,
#     0xFFFFFFFF: Pcaph_mode,
#     0xFFFFFFFE: Pkt_mode,
#     0x00000001: EthII_mode,
#     0x00000002: IPv4_mode,
#     0x00000003: TCP_mode
# }
